#!/bin/env python3

import numpy
import torch

import sdb


def test_tensors(p=None):
    data = [[1, 2],[3, 4]]
    np_array = numpy.array(data)
    x_np = torch.from_numpy(np_array)
    print('From a NumPy array: %r' % x_np)
    x_data = torch.tensor(data)
    print('X: Directly from data: %r' % x_data)
    x_ones = torch.ones_like(x_data) # retains the properties of x_data
    print(f'Ones Tensor like X: {x_ones}')
    x_rand = torch.rand_like(x_data, dtype=torch.float) # overrides the datatype of x_data
    print(f'Random Tensor like X: {x_rand}')
    shape = (6, 4,)
    ones_tensor = torch.ones(shape)
    print(f'Ones Tensor by shape: {ones_tensor}')
    zeros_tensor = torch.zeros(shape)
    print(f'Zeros Tensor by shape: {zeros_tensor}')
    rand_tensor = torch.rand(shape)
    print(f'Random Tensor by shape: {rand_tensor}')

    print('\n === Tensor Attributes:')
    t0 = torch.zeros(shape)
    print(f'\t Shape: {t0.shape}')
    print(f'\t Datatype: {t0.dtype}')
    print(f'\t Device tensor is stored on: {t0.device}')

    print('\n === Tensor Operations')
    print(f'Is cuda available: {torch.cuda.is_available()}')
    print(f'Is vulkan available: {torch.is_vulkan_available()}')

    print(f'Tensor data: {t0}')
    t0[:, 0] = 1.1
    t0[1:, 1] = 0.3
    t0[1:5, 2:3] = 0.91
    t0[len(t0) - 1, len(t0[0]) - 1] = 0.5
    print(f'Standard numpy-like indexing and slicing: {t0}')
    t21 = torch.cat([t0, t0], dim=1)
    print(f'Joining tensors by dimension 1: {t21}')
    t22 = t0 * t0
    print(f'Multiplying tensors: {t22}')
    t23 = t0 @ t0.T
    print(f'Matrix multiplication: {t23}')



def entity(p=None):
    print('PyTorch version: %s (%s)' % (torch.__version__, torch.__path__))
    menu = [
            { 'name': '\n---- Tensors', },
            { 'name': 'Tensors', 'func': test_tensors, },
            ]
    sdb.menu(menu, 'PyTorch test', 2)



if __name__ == '__main__':
    entity()
else:
    pass

